﻿using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Net;
using System.Net.Sockets;
using System.Threading;

namespace PhoneNetTools.Dns
{
    /// <summary>
    /// A generic resolver that is able to send a DNS message to a server using UDP, and return the response to the caller.
    /// Maintains a list of name servers to use.
    /// </summary>
    public class DnsResolver
    {
        /// <summary>
        /// The port used for DNS requests; this is <c>53</c>.
        /// </summary>
        public const int DnsPort = 53;

        /// <summary>
        /// The maximum supported size for DNS requests/responses. This is <c>512</c>.
        /// </summary>
        public const int MaximumDnsUdpMessageSize = 512;

        #region Name servers

        private static readonly Random _random = new Random();
        private static readonly List<IPAddress> _nameServers = new List<IPAddress>();

        /// <summary>
        /// Gets the current collection of name servers all resolver instances use to pick a particular server randomly.
        /// </summary>
        public static ReadOnlyCollection<IPAddress> NameServers
        {
            get
            {
                lock (_nameServers)
                {
                    return _nameServers.AsReadOnly();
                }
            }
        }

        /// <summary>
        /// Clears the current list of name servers.
        /// </summary>
        public static void ClearNameServers()
        {
            lock (_nameServers)
            {
                _nameServers.Clear();
            }
        }

        /// <summary>
        /// Adds an entry to the list of name servers.
        /// </summary>
        /// <param name="address">The address.</param>
        public static void AddNameServer(IPAddress address)
        {
            lock (_nameServers)
            {
                _nameServers.Add(address);
            }
        }

        /// <summary>
        /// Randomly selects an entry from the list of name servers.
        /// </summary>
        /// <returns></returns>
        protected static IPAddress SelectRandomNameServer()
        {
            lock (_nameServers)
            {
                if (_nameServers.Count == 0)
                {
                    return null;
                }

                var index = _random.Next(0, _nameServers.Count);
                var nameServer = _nameServers[index];
                return nameServer;
            }
        }

        #endregion

        /// <summary>
        /// Initializes the <see cref="DnsResolver"/> class and adds the two OpenDNS name servers to the list of used name servers.
        /// </summary>
        static DnsResolver()
        {
            // these are the two OpenDNS name servers
            _nameServers.Add(IPAddress.Parse("208.67.222.222"));
            _nameServers.Add(IPAddress.Parse("208.67.220.220"));
        }

        /// <summary>
        /// Begins a DNS request using the provided message.
        /// </summary>
        /// <param name="message">The DNS message to send to a name server from the name servers list.</param>
        /// <param name="callback">The callback to invoke after the operation has finished.</param>
        /// <param name="state">A user state to identify a particular request.</param>
        /// <returns>An async result.</returns>
        public IAsyncResult BeginDnsRequest(Message message, AsyncCallback callback, object state)
        {
            if (message == null)
            {
                throw new ArgumentNullException("message");
            }

            if (message.Questions.Count < 1)
            {
                throw new ArgumentException("The message contains no questions.");
            }

            // create an AsyncResult<T>
            var ar = new AsyncResult<Message>(callback, state);

            // wrap the required data
            var data = new InternalDataContainer
                           {
                               AsyncResult = ar,
                               Message = message
                           };

            // perform the actual work on the thread pool
            ThreadPool.QueueUserWorkItem(ExecuteDnsRequest, data);
            return ar;
        }

        /// <summary>
        /// Ends an ongoing DNS request operation. Throws any pending exceptions.
        /// </summary>
        /// <param name="asyncResult">The async result returned from the <see cref="BeginDnsRequest"/> method.</param>
        /// <returns>The DNS message received from the server.</returns>
        public Message EndDnsRequest(IAsyncResult asyncResult)
        {
            var ar = asyncResult as AsyncResult<Message>;
            if (ar == null)
            {
                throw new ArgumentException("Invalid IAsyncResult used to called EndDnsRequest.");
            }

            // use the async result to end the operation
            return ar.EndInvoke();
        }

        private static void ExecuteDnsRequest(object state)
        {
            // get back the data
            var data = (InternalDataContainer)state;
            var ar = (AsyncResult<Message>)data.AsyncResult;

            Socket socket = null;
            SocketAsyncEventArgs socketEventArgs = null;
            try
            {
                // create the socket and socket event args
                socket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
                socketEventArgs = new SocketAsyncEventArgs();
                socketEventArgs.UserToken = data;
                socketEventArgs.Completed += SocketEventArgs_Completed;

                // get and set the name server to use
                var ipAddress = SelectRandomNameServer();
                if (ipAddress == null)
                {
                    throw new InvalidOperationException("No name servers given, cannot proceed.");
                }
                socketEventArgs.RemoteEndPoint = new IPEndPoint(ipAddress, DnsPort);

                // set the dns request message
                var buffer = data.Message.ToRawData();
                socketEventArgs.SetBuffer(buffer, 0, buffer.Length);

                // go!
                socket.SendToAsync(socketEventArgs);
            }
            catch (Exception ex)
            {
                // clean up
                if (socketEventArgs != null)
                {
                    socketEventArgs.Dispose();
                }
                if (socket != null)
                {
                    socket.Close();
                }

                // in case of an error, use the exception for the async result
                ar.SetAsCompleted(null, ex, false);
            }
        }

        private static void SocketEventArgs_Completed(object sender, SocketAsyncEventArgs e)
        {
            // get back the data container and socket
            var data = e.UserToken as InternalDataContainer;
            var asyncResult = data.AsyncResult as AsyncResult<Message>;
            var socket = sender as Socket;

            // check for errors first
            if (e.SocketError != SocketError.Success)
            {
                // clean up
                e.Dispose();
                socket.Close();

                // signal the error to the outside
                var exception = new SocketException((int)e.SocketError);
                asyncResult.SetAsCompleted(null, exception, false);
                return;
            }

            try
            {
                // check last operation to see what we need to do
                switch (e.LastOperation)
                {
                    case SocketAsyncOperation.SendTo:
                        // we've sent the datagram successfully, now start receiving
                        // => use the maximum buffer size specified for DNS
                        e.SetBuffer(new byte[MaximumDnsUdpMessageSize], 0, MaximumDnsUdpMessageSize);
                        socket.ReceiveFromAsync(e);
                        break;
                    case SocketAsyncOperation.ReceiveFrom:
                        // we've received an answer from the server
                        var buffer = e.Buffer;

                        // clean up first
                        e.Dispose();
                        socket.Close();

                        // now process response
                        var dnsMessage = new Message();
                        dnsMessage.InitializeFrom(buffer);

                        // signal the result to the outside world
                        asyncResult.SetAsCompleted(dnsMessage, null, false);
                        break;
                    default:
                        var exception = new Exception("Socket event args reported invalid operation: " + e.LastOperation);
                        asyncResult.SetAsCompleted(null, exception, false);
                        break;
                }
            }
            catch (Exception ex)
            {
                // clean up
                e.Dispose();
                socket.Close();

                // signal we're done
                asyncResult.SetAsCompleted(null, ex, false);
            }
        }

        /// <summary>
        /// Just an internal data container to pass around all required information between methods and threads.
        /// </summary>
        private class InternalDataContainer
        {
            public IAsyncResult AsyncResult
            {
                get;
                set;
            }

            public Message Message
            {
                get;
                set;
            }
        }
    }
}
